import numpy as np
import pandas as pd
import scipy.stats as stats
import itertools
from scipy import special

"""最主要的函数只有三个，分别对应者模型的不同层，从高到低分别对应着m_theta(),attribute(),simu_resp"""
class Long_Data_gener(object):
    def __init__(self):
        super(object, self).__init__()

    def m_theta(self, t, mean, growth_rate, n, cor=0.8, input_mean=False, input_cor_mat=None):
        """
        生成被试的时序高阶能力
        :param t: 测量的次数,type=int
        :param mean: 初始能力均值,type=float
        :param growth_rate:测试与测试之间增长的能力均值,type=float
        :param n: 被试数量,type=int
        :param cor: 前后能力的相关,type=float,default=0.8
        :param input_mean: 是否自己输入均值向量,type=arr,default=False
        :param input_cor_mat: 是否自己输入协方差矩阵,type=arr,default=False
        :return: 返回随机生成的被试在不同时间段的能力,type=arr
        """

        if input_cor_mat is None:
            # 根据输入的相关确定协方差矩阵
            # 1. 生成对角矩阵
            eye_mat = np.eye(t)
            # 2.先生成一个所有元素全为协方差系数的矩阵，再在对角线上填入对应的值完成写方法矩阵的建构
            cov_mat = eye_mat - cor + np.ones((t, t)) * cor
        else:
            cov_mat = input_cor_mat

        if input_mean == False:
            # 利用列表生成式生成后一个元素比前一个元素大growth_rate的列表，作为多元正态分布的均值
            mean_arr = [mean + i * growth_rate for i in range(t)]
        else:
            mean_arr = input_mean

        # 利用scipy的多元正态分布函数生成被试的广义能力
        g_theta = stats.multivariate_normal.rvs(mean=mean_arr, cov=cov_mat, size=n)
        return g_theta

    def attribute(self, theta_t, beta_t, delta_t, return_possi=False):
        """
        通过每个时段的高阶能力计算出被试每个时间点的属性
        :param theta_t: 被试的高阶能力他theta,type=arr,shape=(被试个数n,时间点数t)
        :param beta_t: t时间点计算被试属性的斜率，type=arr,shape=(属性个数k，时间点数t)
        :param delta_t: t时间点计算被试属性的截距，type=arr,shape=(属性个数k，时间点数t)
        :return: 不同时间点被试的属性掌握模式,type=arr,shape=(时间点数t,被试个数n,属性个数k)
        """
        att_list = []
        for t in range(theta_t.shape[1]):
            # 因为每一时间段的参数不同，所以分时间点对其进行计算
            # 利用向量的外积进行计算，theta向量和beta向量运算后可获得矩阵
            att_t = np.outer(theta_t[:, t], beta_t) + delta_t
            att_list.append(att_t)
        # 得到没有logistic转化的结果
        x = np.array(att_list)
        # 进行logistic转换
        x = np.exp(x) / (1 + np.exp(x))
        # 将概率转化为具体0或1的属性
        _att_t_list = []
        for t in range(theta_t.shape[1]):
            # 因为dataframe只接受二维数组，因此按照时间点分别转换
            _att_t = pd.DataFrame(x[t, :, :]).applymap(lambda _x_: np.random.binomial(1, _x_, 1)).values.astype(int)
            _att_t_list.append(_att_t)
        # 判断是否返回概率矩阵
        if return_possi == True:
            return np.array(_att_t_list), x
        return np.array(_att_t_list)

    def simu_resp(self, t, alpha, Q, lambda_0, lambda_k, return_possi=False):
        """
        简化后的long-dina
        :param t: 时间点数,type=int
        :param alpha: 被试掌握模式，type=arr,shape=(时间点t,被试个数n，属性个数k)
        :param Q: Q矩阵,type=arr,shape=(时间点t，题目个数i，属性个数k)
        :param lambda_0:截距,type=arr,shape=(时间点数t，题目个数i)/(时间点数t*题目个数i)
        :param lambda_k:斜率,type=arr,shape=(时间点数t，题目个数i)/(时间点数t*题目个数i)
        :param return_possi:是否返回作答为1的概率
        :return:返回模拟的作答矩阵
        """
        # 生成理想作答eta
        eta_list = []
        for t_ in range(t):
            eta_ = self.dina_ideal_resp(Q[t_,:,:],alpha[t_,:,:])
            eta_list.append(list(eta_))
        # 加上截距和斜率,获得没有经过logistic转换的结果
        row_resp = lambda_0.reshape(t, 1, Q.shape[1]) + lambda_k.reshape(t, 1, Q.shape[1]) * np.array(
            eta_list)  # +s_m*gamma_nm_arr.reshape(t,1,Q.shape[1])
        # 进行logistic转换
        possi_resp = np.exp(row_resp) / (1 + np.exp(row_resp))
        resp_list = []
        for t in range(Q.shape[0]):
            # 因为dataframe只接受二维数组，因此按照时间点分别转换
            _resp_t = pd.DataFrame(possi_resp[t, :, :]).applymap(lambda x: np.random.binomial(1, x, 1)).values.astype(
                int)
            resp_list.append(_resp_t)

        if return_possi == True:
            return np.array(resp_list), possi_resp
        return np.array(resp_list)

    def all_pattern(self, k, k_l=2):
        """
        得到所有可能的属性模式
        Args:
            k: 属性的个数
            k_l:属性的水平数

        Returns：
            所有可能的属性
        """
        ite = [i for j in range(k) for i in range(k_l)]
        ite = np.array(ite).reshape(-1, k_l)
        all_pattern = list(itertools.product(*ite))
        return np.array(all_pattern)

    def __normal_I_Q(self, k, k_l):
        """
        隐私函数，生成对角阵作为可达矩阵
        :param k: 属性数量,type=int
        :param k_l: 属性水平数,type=int
        :return: 对角可达矩阵
        """
        # 生成k_l-1个对角矩阵作为可达矩阵
        contrainer = np.vstack([np.eye(k) * (i + 1) for i in range(k_l - 1)])
        return contrainer

    def Q_generate(self, item_cnt, k_cnt, k_l, m_I_mat=True, input_I=None, input_all_pattern=None):
        """
        按照需求生成Q矩阵
        :param item_cnt: 题目数量，type=int
        :param k_cnt: 属性数量,type=int
        :param k_l: 属性水平,type=int
        :param m_I_mat: 是否允许多个可达矩阵,type=bool,default=Ture
        :param input_I: 是否自己输入可达矩阵,type=bool,default=False
        :param input_all_pattern: 是否自己输入所有掌握模式,type=bool,default=None
        :return: 生成的Q矩阵,type=arr
        """
        # 判断是否自己输入所有属性掌握模式矩阵，如果自己输入则直接使用输入矩阵，如果未输入，则调用函数生成。
        if input_all_pattern is not None:
            all_pattern = input_all_pattern
        else:
            all_pattern = self.all_pattern(k_cnt, k_l)
        # 是否自己输入可达矩阵
        if input_I is None:
            I = self.__normal_I_Q(k_cnt, k_l)
        else:
            I = input_I
        # 如果题目长度够长，则使用两个可达矩阵
        if item_cnt >= len(I) * 2 and m_I_mat == True:
            I = np.vstack((I, I))
        x = all_pattern[1:]
        while True:
            # 除去可达矩阵，剩下的题目在all_pattern中随机抽取，判断满足识别条件后，返回arr
            x_ = x[np.random.choice(x.shape[0], item_cnt - len(I))]
            x_list = np.array([(x_[:, i] != x_[:,j]) for i in range(k_cnt) for j in range(k_cnt) if i != j]).sum(axis=1)>0
            if np.product(x_list) == 1:
                break
        container = np.vstack((I, x_))
        return container.astype(np.int8)

    def Q_generate_long(self, t, item_cnt, k_cnt, k_l=2, m_I_mat=True, input_I=None, input_all_pattern=None):
        """重复生成Q矩阵的函数t次"""
        Q_list = []
        for t_ in range(t):
            Q_list.append(self.Q_generate(item_cnt, k_cnt, k_l, m_I_mat, input_I, input_all_pattern))
        return np.array(Q_list)

    def mean_to_lambda(self, s, g):
        """
        将s与g转换为lambda_0和lambda_k
        :param s: 失误参数,type=float
        :param g: 猜测参数,type=float
        :return: 转换后的参数,type=float
        """
        return special.logit(s), special.logit(1 - s) - special.logit(g)

    def dina_ideal_resp(self, Q, att):
        """
        根据DINA模型生成被试的理想作答
        :param Q: Q矩阵,type=arr,shape=(题目个数i,属性个数k)
        :param att: 被试属性,type=arr,shape=(被试个数n，属性个数k)
        :return: DINA模型理想作答,type=arr,shape=(被试个数n，题目个数i)
        """
        eta = np.zeros((len(att), len(Q)))
        for i in range(Q.shape[0]):
            eta[:, i] = np.product(att >= Q[i], axis=1)
        return eta

    def long_dina_ideal_resp(self,t,Q,att):
        """调用同一个函数t次"""
        return np.array([self.dina_ideal_resp(Q[t_,:,:],att) for t_ in range(t)])


